package org.example.processor;

import lombok.extern.slf4j.Slf4j;
import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.VGG16ImagePreProcessor;

import java.io.File;
import java.io.IOException;
@Slf4j
public class ImageProcessor {
    private static final int HEIGHT = 224;
    private static final int WIDTH = 224;
    private static final int CHANNELS = 3;

    public static INDArray preprocessImage(String imagePath)  {
        try {
            NativeImageLoader loader = new NativeImageLoader(HEIGHT, WIDTH, CHANNELS);
            INDArray image = loader.asMatrix(new File(imagePath));
            // 使用VGG16专用预处理器
            new VGG16ImagePreProcessor().transform(image);
            return image;
        } catch (IOException e) {
            log.error("文件问题{}",e);
            throw new RuntimeException(e);
        }
    }

    /**
     * 关键步骤：
     *
     * 使用 model.feedForward() 获取所有层输出
     *
     * 通过层名称从 Map 中提取目标输出
     *
     * 根据需求处理特征形状（展平、池化等）
     *
     * 应用选择：
     *
     * 浅层特征（卷积层）：适合图像局部匹配（如纹理搜索）
     *
     * 深层特征（全连接层）：适合语义级相似度（如商品类别匹配）
     * @param image
     * @param model
     * @return
     */
    public static INDArray extractFeatures(INDArray image, ComputationGraph model) {
        model.setInputs(image);
//        return model.feedForward()[1]; // 获取倒数第二层输出作为特征
        return model.feedForward().get("fc2"); // 获取倒数第二层输出作为特征

    }
}